from args import parse_args
from data.datamodule import FundusDataModule
from models.Vilref import ViLref
from lightning.pytorch.loggers import TensorBoardLogger
from lightning.pytorch.callbacks import ModelCheckpoint
import lightning as L
if __name__ == "__main__":
    args, cfg = parse_args()
    # Load Data Module
    logger = TensorBoardLogger("logs/", name="training")
    data_module = FundusDataModule(
        data_dir=cfg['data']['data_dir'],
        batch_size=cfg['trainer']['batch_size'],
        input_size=cfg['data'].get('input_size', 224),
        num_workers=cfg['trainer'].get('num_workers', 4),
        dataset_name=cfg['data']['dataset_name'],
        split_name=cfg['data']['split_name'],
        use_val=cfg['data'].get('use_val', False)
    )

    # Load Model
    model = ViLref(
        backbone=cfg['model']['backbone'],
        embed_dim=cfg['model']['embed_dim'],
        num_class=cfg['data']['num_classes'],
        opti_cfg=cfg['optimizer'],
        scheduler_cfg=cfg['scheduler'],
        batch_size=cfg['trainer']['batch_size'], 
        loss_terms=cfg['model'].get('loss_terms', None)
        )

    checkpoint_callback = ModelCheckpoint(
        monitor='val_auc',  # 可改为 'val_auc'
        mode='max',
        save_top_k=1,  # 只保存最好的 1 个模型
        verbose=True,  # 显示保存信息
        dirpath='checkpoints/',  # 保存目录
        filename=f"{cfg['data']['dataset_name']}_{cfg['data']['split_name']}"  # 文件名
    )
    # Trainer
    trainer = L.Trainer(
        max_epochs=cfg['trainer']['num_epochs'],
        accelerator=cfg['trainer']['accelerator'],
        precision=cfg['trainer']['precision'],
        logger=logger,
        devices=cfg['trainer']['devices'],
        default_root_dir=cfg['trainer']['save_dir'],
    )

    # Train
    if cfg['data']['use_val']:
        trainer.fit(model, datamodule=data_module)
    else:
        data_module.setup()
        train_loader=data_module.train_dataloader()
        trainer.fit(model, train_dataloaders=train_loader)
    
    # save
    trainer.save_checkpoint(f"{cfg['checkpoints_dir']}/{cfg['data']['dataset_name']}_{cfg['data']['split_name']}.ckpt")